import numpy as np
import math
import librosa
import temporal.util as util


def st_energy(framed_seq):
    energy = np.sum(np.square(framed_seq), 1)  # 求每一帧的L2能量
    return energy


def st_amp(framed_seq):
    amp = np.sum(np.abs(framed_seq), 1)  # 求每一帧的L1能量
    return amp


def st_zcr2(original_seq, frame_length=2048, hop_length=512, center=True):
    zcr = librosa.feature.zero_crossing_rate(original_seq, frame_length=frame_length, hop_length=hop_length,
                                             center=center)
    zcr = np.squeeze(zcr)
    return zcr


def st_zcr(framed_seq):
    lenth, num = framed_seq.shaoe
    zcr = np.zeros((lenth,), dtype=float)
    for i in range(lenth):
        count = 0
        for j in range(num - 1):
            if (framed_seq[i, j] * framed_seq[i, j + 1] < 1):
                count = count + 1
        zcr[i] = float(count)
    return zcr


def double_gate(original_seq, framed_seq, frame_length=2048, hop_length=512, center=False, threshold_high=0,
                threshold_low=-0.05, threshold_zcr=50):
    energy = st_energy(framed_seq)
    amp = st_amp(framed_seq)
    zcr = st_zcr2(original_seq, frame_length=frame_length, hop_length=hop_length, center=center)
    length_e = len(energy)
    length_a = len(amp)
    length_z = len(zcr)
    if (length_a != length_e or length_a != length_z or length_e != length_z):
        print("FALSE！！未对齐")
        return
    energy_N = np.mean(energy)
    energy_sigma = np.var(energy, ddof=1)
    amp_N = np.mean(amp)
    amp_sigma = np.sqrt(np.var(amp, ddof=1))
    energy_normal = (energy - energy_N) / energy_sigma
    amp_normal = (amp - amp_N) / amp_sigma
    # print(energy_normal)
    # print(amp_normal)
    energy_normal = amp_normal

    point_list = []
    mark = 0
    point_list.append(0)
    for i in range(len(energy_normal)):
        if (mark == 0):
            if (energy_normal[i] >= threshold_high):
                point_list.append(i)
                mark = 1
        else:
            if (energy_normal[i] <= threshold_low):
                point_list.append(i)
                mark = 0
    point_list.append(len(energy_normal) - 1)
    begin = point_list[1]
    end = point_list[-2]
    j = begin
    i = end
    while (j >= 0):
        if (zcr[j] < threshold_zcr):
            break
    begin = j
    while (i < length_z):
        if (zcr[i] < threshold_zcr):
            break
    end = i
    return begin, end, energy, amp, zcr


def linear_normalize(data, T=-1):
    """
    如果有给定的归一长度T，那就将长度归一化至T；如果没有给定，则先找到最大长度，令之为T。
    :param data: 10*sample_num*3*frames 的四维list或np.ndarray。
    :param T: 给定的归一的帧数（即长度），调用时可以不给。如果不给T,先找到原数据中最长的frames，以它作为T。
    :return: 10*sample_num*3*T 的四维list。
    """
    if T == -1:
        ten = len(data)
        for num in range(ten):
            sample_numbers = len(data[num])
            for sample in range(sample_numbers):
                three = len(data[num][sample])
                for feature in range(three):
                    frames = len(data[num][sample][feature])
                    T = max(frames, T)

    normalized_data = []
    ten = len(data)
    for num in range(ten):
        num_list = []
        sample_numbers = len(data[num])
        for sample in range(sample_numbers):
            sample_list = []
            three, frames = np.shape(data[num][sample])
            for feature in range(three):
                feature_list = []
                seq = data[num][sample][feature]
                for new_frame in range(T):
                    if new_frame != T - 1:
                        frame = new_frame / (T - 1) * (frames - 1)
                        frame1 = math.floor(frame)
                        frame2 = frame1 + 1
                        w1 = frame2 - frame
                        w2 = frame - frame1
                        feature_list.append(w1 * seq[frame1] + w2 * seq[frame2])
                    else:
                        feature_list.append(seq[frames - 1])
                sample_list.append(feature_list)
            num_list.append(sample_list)
        normalized_data.append(num_list)

    len_ten = []
    len_already = [0]
    len_total = 0
    for num in range(ten):
        len_ten.append(len(normalized_data[num]))
        len_total += len_ten[num]
        len_already.append(len_total)
    reshaped_array = np.zeros((len_total, three * T), dtype='float64')
    for i in range(ten):
        for j in range(len_ten[i]):
            reshaped_array[len_already[i] + j] = np.array(normalized_data[i][j]).reshape(-1)
    label_array = np.zeros(len_total, dtype='int32')
    for i in range(len(label_array)):
        index = 9
        for j in range(ten - 1):
            if len_already[j + 1] >= i:
                index = j
                break
        label_array[i] = index

    # array = np.array(normalized_data)
    # reshaped_array = np.zeros((ten * sample_numbers, three * T), dtype='float64')
    # for i in range(ten):
    #     for j in range(sample_numbers):
    #         reshaped_array[i * sample_numbers + j] = array[i, j, :, :].reshape(-1)
    # label_array = np.zeros(ten * sample_numbers, dtype='int32')
    # for i in range(len(label_array)):
    #     label_array[i] = i // sample_numbers
    return reshaped_array, label_array, T


def prefeature(all_data, wrap_data):
    cls_num = wrap_data.shape[0]
    data_num = wrap_data.shape[1]
    prefeature = []
    for i in range(cls_num):
        prefeature.append([])
        for j in range(data_num):
            tmp = []
            begin, end, ener, amp, zcr = double_gate(all_data[i][j], wrap_data[i][j], frame_length=882, hop_length=441)
            tmp.append(ener[begin: end + 1].tolist())
            tmp.append(amp[begin: end + 1].tolist())
            tmp.append(zcr[begin: end + 1].tolist())
            if np.array(tmp).shape[1] > 10:
                prefeature[i].append(tmp)
            # print(np.array(tmp).shape)

    return prefeature


def single_audio(array, T, sr=44100, frame_length=882, shift_length=441, window_function="juxing", hop_length=512,
                 center=True):
    wrapped_array = util.devideAndWindow(array, sr, frame_length, shift_length, window_function)
    prefeature_data = []

    print(np.array(wrapped_array).shape)
    # print(wrapped_array)
    begin, end, ener, amp, zcr = double_gate(array, wrapped_array, frame_length=882, hop_length=441)

    prefeature_data.append(ener[begin: end + 1].tolist())
    prefeature_data.append(amp[begin: end + 1].tolist())
    prefeature_data.append(zcr[begin: end + 1].tolist())
    X_0 = np.ndarray((3 * T, 1), dtype='float32')
    prefeature_array = np.array(prefeature_data)
    frames = prefeature_array.shape[1]

    sample_list = []
    three, frames = np.shape(prefeature_array)
    for feature in range(three):
        feature_list = []
        seq = prefeature_array[feature]
        for new_frame in range(T):
            if new_frame != T - 1:
                frame = new_frame / (T - 1) * (frames - 1)
                frame1 = math.floor(frame)
                frame2 = frame1 + 1
                w1 = frame2 - frame
                w2 = frame - frame1
                feature_list.append(w1 * seq[frame1] + w2 * seq[frame2])
            else:
                feature_list.append(seq[frames - 1])
        sample_list.append(feature_list)

    sample_array = np.array(sample_list)
    print(sample_array.shape)
    print(X_0.shape)
    X_0 = np.concatenate((sample_array[0], sample_array[1], sample_array[2]), axis=0).reshape(1,-1)
    print(X_0.shape)
    return X_0


if __name__ == '__main__':
    wrapped_data, all_data = util.activateFromOutside(path="./data", sr=44100, frame_length=0.02, shift_length=0.01,
                                                      window_function="juxing")
    prefeature = prefeature(all_data, wrapped_data)
